import { useState, useRef, useEffect } from 'react';
import { marked } from 'marked';
import hljs from 'highlight.js';

// 配置 marked 使用 highlight.js 高亮代码
marked.use({
    renderer: {
        code(info) {
            const validLanguage = info.lang || 'plaintext';
            return `<pre><code class="hljs ${validLanguage}">${hljs.highlight(info.text, { language: validLanguage }).value}</code></pre>`;
        },
    },
    gfm: true,
    breaks: true,
})

interface Message {
  role: 'user' | 'assistant' | 'system';
  content: string;
}

export default function OllamaChat() {
  // 状态管理
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState('');
  const [isLoading, setIsLoading] = useState(false);
  const [systemPrompt, setSystemPrompt] = useState(
    '你是一个乐于助人的AI助手，请用中文回答，保持回答简洁专业。'
  );
  const [showPromptEditor, setShowPromptEditor] = useState(false);
  const abortControllerRef = useRef<AbortController | null>(null);
  const messagesEndRef = useRef<HTMLDivElement>(null);

  // 初始化：从 localStorage 加载历史记录和提示词
  useEffect(() => {
    const savedMessages = localStorage.getItem('ollama_chat_messages');
    const savedPrompt = localStorage.getItem('ollama_system_prompt');
    if (savedMessages) setMessages(JSON.parse(savedMessages));
    if (savedPrompt) setSystemPrompt(savedPrompt);
    return function () {
        // 清理函数：在组件卸载时取消请求
        abortControllerRef.current?.abort();
    }
  }, []);

  // 持久化存储
  useEffect(() => {
    localStorage.setItem('ollama_chat_messages', JSON.stringify(messages.slice(-200))); // 限制存储数量
    // 仅在提示词变化时更新 localStorage，避免频繁写入
    localStorage.setItem('ollama_system_prompt', systemPrompt);
  }, [messages, systemPrompt]);

  // 自动滚动
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

  // 发送消息（包含系统提示词）
  const sendMessage = async () => {
    if (!input.trim() || isLoading) return;

    const userMessage: Message = { role: 'user', content: input };
    const newMessages = [...messages, userMessage].slice(-20); // 限制消息数量，避免过长的历史记录
    setMessages(newMessages);
    setInput('');
    setIsLoading(true);

    abortControllerRef.current = new AbortController();

    try {
      // 构造包含系统提示的消息数组
      const messagesWithSystem = [
        { role: 'system', content: systemPrompt },
        ...newMessages.filter(m => m.role !== 'system'), // 避免重复系统消息
      ];

      const response = await fetch('/api/chat', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: 'qwen2',
          messages: messagesWithSystem,
          options: { num_ctx: 256 },
          stream: true,
        }),
        signal: abortControllerRef.current.signal,
      });

      if (!response.ok || !response.body) {
        throw new Error(`请求失败: ${response.status}`);
      }

      // 添加初始 assistant 消息
      setMessages(prev => [...prev, { role: 'assistant', content: '' }]);

      const reader = response.body.getReader();
      const decoder = new TextDecoder();

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        const chunk = decoder.decode(value);
        chunk.split('\n').forEach(line => {
          if (!line.trim()) return;
          try {
            const data = JSON.parse(line);
            if (data.message?.content) {
              setMessages(prev => {
                const last = prev[prev.length - 1];
                return last.role === 'assistant'
                  ? [...prev.slice(0, -1), { ...last, content: last.content + data.message.content }]
                  : prev;
              });
            }
          } catch (e) {
            console.error('解析错误:', e);
          }
        });
      }
    } catch (e: any) {
      if (e.name !== 'AbortError') {
        setMessages(prev => [...prev, { role: 'assistant', content: `错误: ${e.message}` }]);
      }
    } finally {
      setIsLoading(false);
      abortControllerRef.current = null;
    }
  };

  // 停止生成
  const stopGenerating = () => {
    abortControllerRef.current?.abort();
    setIsLoading(false);
  };

  // 清空对话（保留系统提示）
  const clearConversation = () => {
    setMessages([]);
  };

  return (
    <div className="flex flex-col h-screen mx-auto p-4 bg-gray-50">
      {/* 顶部工具栏 */}
      <div className="flex justify-between items-center mb-4">
        <h1 className="text-xl font-bold">Ollama Chat</h1>
        <div className="flex space-x-2">
          <button
            onClick={() => setShowPromptEditor(!showPromptEditor)}
            className="px-3 py-1 bg-gray-200 rounded hover:bg-gray-300 text-sm"
          >
            {showPromptEditor ? '隐藏提示词' : '编辑提示词'}
          </button>
          <button
            onClick={clearConversation}
            className="px-3 py-1 bg-gray-200 rounded hover:bg-gray-300 text-sm"
            disabled={isLoading}
          >
            清空对话
          </button>
        </div>
      </div>

      {/* 系统提示词编辑器 */}
      {showPromptEditor && (
        <div className="mb-4 p-3 bg-yellow-50 rounded-lg border border-yellow-200">
          <label className="block text-sm font-medium mb-1">系统提示词：</label>
          <textarea
            value={systemPrompt}
            onChange={(e) => setSystemPrompt(e.target.value)}
            className="w-full p-2 border rounded"
            rows={3}
            placeholder="指导AI行为的系统级提示词..."
          />
          <div className="text-xs text-gray-500 mt-1">
            提示词会作为隐藏消息附加在每次对话前
          </div>
        </div>
      )}

      {/* 聊天消息区域 */}
      <div className="flex-1 overflow-y-auto mb-4 space-y-4">
        {messages.map((msg, index) => (
          <div
            key={index}
            className={`p-3 rounded-lg max-w-[80%] ${
              msg.role === 'user'
                ? 'bg-blue-100 ml-auto'
                : msg.role === 'assistant'
                ? 'bg-green-100 mr-auto'
                : 'bg-yellow-100 mr-auto text-xs italic'
            }`}
          >
            <div className="font-semibold">
              {msg.role === 'user' ? 'You' : msg.role === 'assistant' ? 'AI' : 'System'}
            </div>
            <div className="overflow-x-auto ai-content" dangerouslySetInnerHTML={{
                __html: marked(msg.content),
            }}/>
          </div>
        ))}
        <div ref={messagesEndRef} />
      </div>

      {/* 输入区域 */}
      <div className="flex space-x-2">
        <textarea
          value={input}
          onChange={(e) => setInput(e.target.value)}
          onKeyDown={(e) => e.key === 'Enter' && !e.shiftKey && sendMessage()}
          placeholder="输入消息..."
          className="flex-1 p-2 border rounded-lg resize-none"
          rows={3}
          disabled={isLoading}
        />
        <div className="flex flex-col space-y-2">
          <button
            onClick={sendMessage}
            disabled={!input.trim() || isLoading}
            className="px-4 py-2 bg-blue-500 text-white rounded-lg hover:bg-blue-600 disabled:opacity-50"
          >
            发送
          </button>
          {isLoading && (
            <button
              onClick={stopGenerating}
              className="px-4 py-2 bg-red-500 text-white rounded-lg hover:bg-red-600"
            >
              停止
            </button>
          )}
        </div>
      </div>
    </div>
  );
}